import src.corpus as cs
import numpy as np

# 全局变量区
docs = []      # 二维列表, documents[i][j]表示第i篇文档的第j个单词
K = 9           # 主题个数
V = 0           # 词汇表中词汇的个数
M = 0           # 语料库里文档的个数
alpha = 2.0     # doc-topic先验分布的狄利克雷参数
beta = 0.5      # topic-word先验分布的狄利克雷参数
z = []          # 二维列表, z[i][j]表示第i篇文档的第j个单词的主题编号
n_wt = []     # 二维列表, n_wt[i][j]表示词语i分给主题j的次数
n_dt = []      # 二维列表, n_dt[i][j]表示第i篇文档里归入主题j的词语的个数
n_t = []        # 一维列表, n_t[i]表示第i个主题下词语的个数
n_d = []        # 一维列表, n_d[i]表示第i篇文档中词语的个数
# 二维列表, theta[i][j]表示第i篇文档中主题j的概率;
# theta是doc-topic分布, 它是一个D维向量, D是文档集中文档的个数;
# theta向量中每一个分量代表了一篇文档在K个主题下的分布概率情况;
theta = []
# 二维列表, phi[i][j]表示第i个主题下单词j的概率, 或者说在主题i下所占的权重更好理解些;
# phi是topic-word分布, 它是一个K维向量, K是主题的个数;
# phi向量中每一个分量代表了该主题下单词的分布概率情况, 或者说分布权重;
phi = []
n_theta = []
n_phi = []
n_sum = 0   # 样本的容量
n_iterate = 1000        # 迭代次数


def initial_state(docs):
	z = [[] for _ in range(0, M)]
	n_theta = [[0 for _ in range(0, K)] for _ in range(0, M)]
	n_phi = [[0 for _ in range(0, V)] for _ in range(0, K)]
	# 初始化一堆计数器
	n_wt = [[0 for _ in range(0, K)] for _ in range(0, V)]
	n_dt = [[0 for _ in range(0, K)] for _ in range(0, M)]
	n_t  = [0 for _ in range(0, K)]
	n_d  = [0 for _ in range(0, M)]

	# 为每一篇文档的每个单词都随机赋一个主题
	for m in range(0, M):
		N = len(docs[m])
		z[m] = [0 for _ in range(0, N)]
		for w in range(0, N):
			topic = int(np.random.random() * K)      # 为当前单词随机赋值一个主题编号
			z[m][w] = topic
			n_wt[docs[m][w]][topic] += 1        # 当前单词分配给指定主题的次数+1
			n_dt[m][topic] +=1                  # 第m篇文档中分配给主题topic的次数+1
			n_t[topic] += 1                            # topic主题下词语的数量+1
			n_d[m] += 1                            # 第m篇文档中词语的数量
	return z, n_wt, n_dt, n_t, n_d, n_theta, n_phi

def gibbs():
	for i in range(0, n_iterate):
		print("迭代第%d次" % (i))
		for m in range(0, len(z)):      # 第m篇文档
			for w in range(0, len(z[m])):       # 第m篇文档的第w个单词
				topic = sample_conditional(m, w)
				z[m][w] = topic             # 为m篇文档的第w个单词重新填充主题
		update_params()

def sample_conditional(m, w):
	# 取出第m篇文档第w个单词在上一轮迭代中被分配的主题编号
	topic = z[m][w]
	# 在各计数器中删除掉本单词带来的影响
	# 每次采样只能移动一个维度, 其他维度都作为条件
	# 这里的维度应该指的是一个单词;
	# 毕竟一篇文档算是一个样本, 文档下的所有单词是样本的所有特征, 或者说是样本向量的所有分量;
	n_wt[docs[m][w]][topic] -=1
	n_dt[m][topic] -=1
	n_t[topic] -=1
	n_d[m] -=1

	p = [0 for _ in range(0, K)]
	for k in range(0, K):
		e1 = (n_dt[m][k] + alpha) / (n_d[m] + K * alpha)
		e2 = (n_wt[docs[m][w]][k] + beta) / (n_t[k] + V * beta)
		p[k] = e1 * e2
	for k in range(1, K):
		p[k] += p[k-1]
	# 在一个doc的topic分布被估计（根据gibbs采样公式）完毕之后，算法后续将基于算出来的这个topic概率，
	# 从K个主题中随机选择一个topic，来作为本次迭代中当前doc的topic
	# 原来下面代码是要伪造随机啊，困惑了我这么长时间。。。。。。。。
	# 不过这种伪随机的方式还挺特别的，不知道是不是有什么说法？
	# 以前光注意代码了，没有彻底搞懂LDA的原理。。。。
	u = np.random.random() * p[K-1]
	for topic in range(0, K):
		if (u < p[topic]):
			break
	# 将重新估计的本词语加入计数器组
	n_wt[docs[m][w]][topic] += 1
	n_dt[m][topic] += 1
	n_t[topic] += 1
	n_d[m] += 1
	# 并返回新采样到的主题
	return topic

# 更新phi和theta各自的累积量
def update_params():
	for m in range(0, len(docs)):
		for k in range(0, K):
			n_theta[m][k] += (n_dt[m][k] + alpha) / (n_d[m] + K * alpha)
	for k in range(0, K):
		for v in range(0, V):
			n_phi[k][v] += (n_wt[v][k] + beta) / (n_t[k] + V * beta)

# 计算各主题权重
def count_weight(row):
	topic = ['IT', '体育', '健康', '军事', '招聘', '教育', '文化', '旅游', '财经']
	count = [0 for _ in range(0, len(topic))]
	for i in range(0, len(topic)):
		for j in range(0, len(row)):
			if str(row[j]).find(topic[i]) != -1:
				count[i] +=1
	return count


if __name__ == "__main__":
	# 装载语料库
	cs.load_documents("../data/sample")
	print(cs.corpus)
	# 测试: 输出语料库
	for m in range(0, len(cs.corpus)):
		print(cs.docs_name[m])
		print(str(cs.corpus[m]))
	# 测试: 将语料库输入至文件
	with open('corpus_test.txt', "w") as writer:
		for m in range(0, len(cs.corpus)):
			for w in range(0, len(cs.corpus[m])):
				writer.write(str(cs.corpus[m][w]))
				writer.write('\t')
			writer.write('\n')
	# 测试: 将语料库对应的文档名输入至文件
	with open('doc_name_test.txt', 'w') as writer:
		for i in range(0, len(cs.docs_name)):
			writer.write(cs.docs_name[i])
			writer.write('\n')
	# 将语料库转换成id矩阵
	cs.translate_word_to_id()
	print('finish')

	docs = cs.corpus
	V = len(cs.vocabulary)
	M = len(cs.corpus)
	# 初始化gibbs采样状态
	(z, n_wt, n_dt, n_t, n_d, n_theta, n_phi) = initial_state(cs.corpus)
	# print(len(n_theta))
	# 进行吉布斯采样
	gibbs()

	print(n_theta)
	print(n_phi)
	# 测试: 将theta矩阵输出至文件
	with open('test.txt', 'w') as writer:
		for m in range(0, len(n_theta)):
			for t in range(0, len(n_theta[m])):
				writer.write(str(n_theta[m][t]))
				writer.write('\t')
			writer.write('\n')
	docs_topic = [[] for _ in range(0, K)]
	for m in range(0, len(n_theta)):
		# 取出每个文档里权重最大的一个topic索引
		max_index = np.where(n_theta[m] == np.max(n_theta[m]))[0][0]
		docs_topic[max_index].append(cs.docs_name[m])
	# 测试: 将分好的topic, 按行输出全部文档
	# with open('topic.txt', 'w') as writer:
	# 	for t in range(0, len(docs_topic)):
	# 		for m in range(0, len(docs_topic[t])):
	# 			writer.write(str(docs_topic[t][m]))
	# 			writer.write('\t')
	# 		writer.write('\n')
	counts = []
	for i in range(0, len(docs_topic)):
		counts.append(count_weight(docs_topic[i]))

	topic = ['IT', '体育', '健康', '军事', '招聘', '教育', '文化', '旅游', '财经']
	with open('topic.txt', 'w') as writer:
		for t in range(0, len(counts)):
			for m in range(0, len(counts[t])):
				writer.write(topic[m] + ": " + str(counts[t][m]))
				writer.write('\t')
			writer.write('\n')

